import streamlit as st
import sys
import os
from typing import Dict, Any
from datetime import datetime
import json
from loguru import logger

# 添加当前目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

# 导入agent相关模块
from flood_forecast_agent import FloodForecastAgent
try:
    from __init__ import OLLAMA_LLM_CONFIG
except ImportError:
    # 如果导入失败，使用默认配置
    OLLAMA_LLM_CONFIG = {
        "model": "qwq:latest",
        "type": "ollama",
        "base_url": "http://10.48.0.81:11434",
        "temperature": 0.1
    }

# 配置loguru日志
logger.remove()  # 移除默认配置
logger.add(
    sys.stdout,
    level="INFO",
    format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>",
    colorize=True
)

# 页面配置
st.set_page_config(
    page_title="洪水预报智能代理",
    page_icon="🌊",
    layout="wide",
    initial_sidebar_state="expanded"
)

# 自定义CSS样式
st.markdown("""
<style>
.main-header {
    background: linear-gradient(90deg, #1e3c72 0%, #2a5298 100%);
    padding: 1rem;
    border-radius: 10px;
    color: white;
    text-align: center;
    margin-bottom: 2rem;
}

.status-card {
    background: #f8f9fa;
    padding: 1rem;
    border-radius: 8px;
    border-left: 4px solid #007bff;
    margin: 1rem 0;
}

.parameter-card {
    background: #e8f5e8;
    padding: 1rem;
    border-radius: 8px;
    border-left: 4px solid #28a745;
    margin: 0.5rem 0;
}

.missing-parameter {
    background: #fff3cd;
    padding: 0.5rem;
    border-radius: 5px;
    border-left: 3px solid #ffc107;
    margin: 0.3rem 0;
}

.chat-message {
    padding: 1rem;
    margin: 0.5rem 0;
    border-radius: 10px;
}

.user-message {
    background: #e3f2fd;
    border-left: 4px solid #2196f3;
}

.assistant-message {
    background: #f3e5f5;
    border-left: 4px solid #9c27b0;
}
</style>
""", unsafe_allow_html=True)

def init_session_state():
    """初始化session state"""
    if 'agent' not in st.session_state:
        try:
            st.session_state.agent = FloodForecastAgent(OLLAMA_LLM_CONFIG)
            logger.info("Agent初始化成功")
        except Exception as e:
            logger.error(f"Agent初始化失败: {e}")
            st.error(f"Agent初始化失败: {e}")
            st.stop()
    
    if 'chat_history' not in st.session_state:
        st.session_state.chat_history = []
    
    if 'last_response' not in st.session_state:
        st.session_state.last_response = None
    
    if 'is_processing' not in st.session_state:
        st.session_state.is_processing = False

def display_agent_status(response: Dict[str, Any]):
    """显示agent状态信息"""
    col1, col2, col3 = st.columns(3)
    
    with col1:
        stage_color = {
            "collection": "🔍",
            "validation": "✅", 
            "confirmation": "📋"
        }
        st.markdown(f"""
        <div class="status-card">
            <h4>{stage_color.get(response['current_stage'], '❓')} 当前阶段</h4>
            <p><strong>{response['current_stage'].upper()}</strong></p>
        </div>
        """, unsafe_allow_html=True)
    
    with col2:
        completion_icon = "✅" if response['is_complete'] else "⏳"
        st.markdown(f"""
        <div class="status-card">
            <h4>{completion_icon} 完成状态</h4>
            <p><strong>{'已完成' if response['is_complete'] else '进行中'}</strong></p>
        </div>
        """, unsafe_allow_html=True)
    
    with col3:
        confidence = response.get('confidence_score', 0)
        confidence_color = "#28a745" if confidence > 0.8 else "#ffc107" if confidence > 0.5 else "#dc3545"
        st.markdown(f"""
        <div class="status-card">
            <h4>📊 置信度</h4>
            <p style="color: {confidence_color}"><strong>{confidence:.1%}</strong></p>
        </div>
        """, unsafe_allow_html=True)

def display_parameters(response: Dict[str, Any]):
    """显示参数信息"""
    st.subheader("📋 参数收集状态")
    
    # 已收集参数
    if response['collected_parameters']:
        st.markdown("**✅ 已收集参数:**")
        for key, value in response['collected_parameters'].items():
            st.markdown(f"""
            <div class="parameter-card">
                <strong>{key}:</strong> {value}
            </div>
            """, unsafe_allow_html=True)
    
    # 缺失参数
    if response['missing_parameters']:
        st.markdown("**⚠️ 缺失参数:**")
        for param in response['missing_parameters']:
            st.markdown(f"""
            <div class="missing-parameter">
                {param}
            </div>
            """, unsafe_allow_html=True)
    
    # 验证结果
    if response.get('validation_result'):
        st.markdown("**🔍 验证结果:**")
        validation = response['validation_result']
        if validation.get('is_valid'):
            st.success("✅ 参数验证通过")
        else:
            st.warning("⚠️ 参数验证存在问题")
            if validation.get('errors'):
                for error in validation['errors']:
                    st.error(f"❌ {error}")
    
    # 如果会话完成且参数验证成功，显示完整的JSON参数列表
    if (response.get('is_complete') and 
        response.get('validation_result', {}).get('is_valid') and 
        response.get('collected_parameters')):
        
        st.markdown("---")
        st.markdown("**🎉 参数收集完成！**")
        st.markdown("**📄 完整参数JSON:**")
        
        # 创建一个可展开的区域显示JSON
        with st.expander("查看完整参数JSON", expanded=True):
            # 格式化JSON显示
            json_str = json.dumps(
                response['collected_parameters'], 
                ensure_ascii=False, 
                indent=2
            )
            
            # 使用代码块显示JSON
            st.code(json_str, language='json')
            
            # 添加复制按钮功能
            st.markdown("**💾 操作:**")
            col1, col2 = st.columns(2)
            
            with col1:
                # 下载按钮
                st.download_button(
                    label="📥 下载JSON文件",
                    data=json_str,
                    file_name=f"flood_forecast_params_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
                    mime="application/json",
                    use_container_width=True
                )
            
            with col2:
                # 复制到剪贴板的提示
                if st.button("📋 复制JSON", use_container_width=True):
                    # 使用JavaScript复制到剪贴板
                    st.markdown(f"""
                    <script>
                    navigator.clipboard.writeText(`{json_str}`);
                    </script>
                    """, unsafe_allow_html=True)
                    st.success("JSON已复制到剪贴板！")
        
        # 显示参数摘要统计
        st.markdown("**📊 参数摘要:**")
        param_count = len(response['collected_parameters'])
        st.info(f"共收集到 {param_count} 个参数")
        
        # 按类别显示参数
        st.markdown("**🏷️ 参数分类:**")
        
        # 基本信息
        basic_params = ['reservoir_name', 'forecast_model', 'target_variable']
        basic_collected = {k: v for k, v in response['collected_parameters'].items() if k in basic_params}
        if basic_collected:
            st.markdown("*基本信息:*")
            for k, v in basic_collected.items():
                st.markdown(f"  • {k}: {v}")
        
        # 时间参数
        time_params = ['start_time', 'forecast_duration', 'rainfall_forecast_period']
        time_collected = {k: v for k, v in response['collected_parameters'].items() if k in time_params}
        if time_collected:
            st.markdown("*时间参数:*")
            for k, v in time_collected.items():
                st.markdown(f"  • {k}: {v}")
        
        # 其他参数
        other_params = {k: v for k, v in response['collected_parameters'].items() 
                       if k not in basic_params + time_params}
        if other_params:
            st.markdown("*其他参数:*")
            for k, v in other_params.items():
                st.markdown(f"  • {k}: {v}")

def display_chat_history():
    """显示聊天历史"""
    st.subheader("💬 对话历史")
    
    if not st.session_state.chat_history:
        st.info("👋 欢迎使用洪水预报智能代理！请开始您的对话。")
        return
    
    for i, message in enumerate(st.session_state.chat_history):
        if message['role'] == 'user':
            st.markdown(f"""
            <div class="chat-message user-message">
                <strong>👤 用户:</strong><br>
                {message['content']}
            </div>
            """, unsafe_allow_html=True)
        else:
            st.markdown(f"""
            <div class="chat-message assistant-message">
                <strong>🤖 助手:</strong><br>
                {message['content']}
            </div>
            """, unsafe_allow_html=True)

def main():
    """主函数"""
    init_session_state()
    
    # 主标题
    st.markdown("""
    <div class="main-header">
        <h1>🌊 洪水预报智能代理</h1>
        <p>基于大语言模型的智能参数收集与验证系统</p>
    </div>
    """, unsafe_allow_html=True)
    
    # 侧边栏 - 系统信息
    with st.sidebar:
        st.header("🔧 系统信息")
        
        # 模型配置
        st.subheader("🤖 模型配置")
        config = OLLAMA_LLM_CONFIG
        st.write(f"**模型:** {config.get('model', 'Unknown')}")
        st.write(f"**类型:** {config.get('type', 'Unknown')}")
        st.write(f"**服务地址:** {config.get('base_url', 'Unknown')}")
        st.write(f"**温度:** {config.get('temperature', 'Unknown')}")
        
        st.divider()
        
        # 操作按钮
        st.subheader("🛠️ 操作")
        
        if st.button("🔄 重置会话", type="secondary", use_container_width=True, disabled=st.session_state.is_processing):
            st.session_state.agent.reset_session()
            st.session_state.chat_history = []
            st.session_state.last_response = None
            st.session_state.is_processing = False  # 重置处理状态
            st.success("会话已重置")
            st.rerun()
        
        # 完成会话按钮（仅在确认阶段显示）
        if (st.session_state.last_response and 
            st.session_state.last_response.get('current_stage') == 'confirmation' and
            not st.session_state.agent.is_session_completed):
            if st.button("✅ 完成会话", type="primary", use_container_width=True, disabled=st.session_state.is_processing):
                st.session_state.agent.mark_session_completed()
                st.success("会话已完成！")
                st.rerun()
        
        if st.button("📊 会话摘要", use_container_width=True, disabled=st.session_state.is_processing):
            summary = st.session_state.agent.get_session_summary()
            st.json(summary)
        
        # 导出功能
        if st.session_state.last_response and st.session_state.last_response.get('is_complete'):
            st.subheader("📤 导出")
            
            # 导出参数JSON
            if st.button("📄 导出参数JSON", use_container_width=True, disabled=st.session_state.is_processing):
                params_json = json.dumps(
                    st.session_state.last_response['collected_parameters'], 
                    ensure_ascii=False, 
                    indent=2
                )
                st.download_button(
                    label="💾 下载参数文件",
                    data=params_json,
                    file_name=f"flood_forecast_params_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
                    mime="application/json",
                    use_container_width=True,
                    disabled=st.session_state.is_processing
                )
    
    # 主内容区域
    col1, col2 = st.columns([2, 1])
    
    with col1:
        # 显示聊天历史
        display_chat_history()
        
        # 用户输入
        st.subheader("✍️ 输入您的需求")
        
        # 快速输入示例
        st.markdown("**💡 示例输入:**")
        example_inputs = [
            "我需要对大伙房水库进行洪水预报",
            "使用新安江模型，从明天上午8点开始",
            "预报48小时，降雨预见期12小时，目标变量是流量"
        ]
        
        cols = st.columns(len(example_inputs))
        for i, example in enumerate(example_inputs):
            with cols[i]:
                if st.button(f"📝 {example[:10]}...", key=f"example_{i}", use_container_width=True, disabled=st.session_state.is_processing):
                    st.session_state.user_input = example
        
        # 文本输入
        user_input = st.text_area(
            "请输入您的洪水预报需求:",
            value=st.session_state.get('user_input', ''),
            height=100,
            placeholder="例如：我需要对某水库进行洪水预报，使用某某模型..."
        )
        
        # 发送按钮
        col_send, col_clear = st.columns([1, 1])
        
        with col_send:
            # 根据处理状态决定按钮是否禁用
            button_disabled = st.session_state.is_processing
            button_text = "⏳ 处理中..." if button_disabled else "🚀 发送"
            
            if st.button(button_text, type="primary", use_container_width=True, disabled=button_disabled):
                if user_input.strip():
                    # 设置处理状态
                    st.session_state.is_processing = True
                    
                    # 添加用户消息到历史
                    st.session_state.chat_history.append({
                        'role': 'user',
                        'content': user_input
                    })
                    
                    # 处理用户输入
                    with st.spinner("🤖 AI正在分析您的需求..."):
                        try:
                            response = st.session_state.agent.process_user_input(user_input)
                            
                            # 添加AI回复到历史
                            st.session_state.chat_history.append({
                                'role': 'assistant',
                                'content': response['response']
                            })
                            
                            # 保存最新响应
                            st.session_state.last_response = response
                            
                            # 清空输入
                            st.session_state.user_input = ''
                            
                        except Exception as e:
                            logger.error(f"处理用户输入时出错: {e}")
                            st.error(f"处理请求时出错: {e}")
                        finally:
                            # 无论成功还是失败，都要重置处理状态
                            st.session_state.is_processing = False
                            
                        st.rerun()
                else:
                    st.warning("请输入您的需求")
        
        with col_clear:
            if st.button("🗑️ 清空输入", use_container_width=True, disabled=st.session_state.is_processing):
                st.session_state.user_input = ''
                st.rerun()
    
    with col2:
        # 显示状态和参数信息
        if st.session_state.last_response:
            display_agent_status(st.session_state.last_response)
            st.divider()
            display_parameters(st.session_state.last_response)
        else:
            st.info("🎯 开始对话后，这里将显示参数收集状态")
    
    # 页脚
    st.markdown("---")
    st.markdown("""
    <div style="text-align: center; color: #666; padding: 1rem;">
        <p>🌊 洪水预报智能代理 | 基于 LangChain & Streamlit 构建</p>
    </div>
    """, unsafe_allow_html=True)

if __name__ == "__main__":
    main()